// Copyright (c) 2024 Go-Frame-Lite
// Go-Frame-Lite is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package middleware

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"strings"

	"go-frame-lite/pkg/crypto"
	"go-frame-lite/pkg/logger"

	"github.com/gin-gonic/gin"
	"github.com/spf13/viper"
	"go.uber.org/zap"
)

// SM3IntegrityConfig SM3完整性校验配置
type SM3IntegrityConfig struct {
	Enabled         bool     // 是否启用
	HashHeader      string   // 哈希头字段名称
	ExcludePaths    []string // 排除的路径
	HashFormat      string   // 哈希格式，支持"hex"或"base64"
	VerifyResponses bool     // 是否验证响应完整性
}

// DefaultSM3IntegrityConfig 默认SM3完整性校验配置
func DefaultSM3IntegrityConfig() SM3IntegrityConfig {
	return SM3IntegrityConfig{
		Enabled:         viper.GetBool("crypto.sm3.enabled"),
		HashHeader:      viper.GetString("crypto.sm3.hash_header"),
		ExcludePaths:    []string{"/health", "/metrics"},
		HashFormat:      viper.GetString("crypto.sm3.hash_format"),
		VerifyResponses: viper.GetBool("crypto.sm3.verify_responses"),
	}
}

// sm3ResponseWriter 自定义响应写入器，用于计算响应体的SM3哈希
type sm3ResponseWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// Write 重写Write方法，将响应体写入缓冲区
func (w sm3ResponseWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// SM3IntegrityMiddleware SM3完整性校验中间件
func SM3IntegrityMiddleware(config ...SM3IntegrityConfig) gin.HandlerFunc {
	cfg := DefaultSM3IntegrityConfig()
	if len(config) > 0 {
		cfg = config[0]
	}

	// 设置默认哈希头名称（如果未配置）
	if cfg.HashHeader == "" {
		cfg.HashHeader = "X-SM3-Hash"
	}

	// 设置默认哈希格式（如果未配置）
	if cfg.HashFormat != "hex" && cfg.HashFormat != "base64" {
		cfg.HashFormat = "hex"
	}

	return func(c *gin.Context) {
		// 检查是否启用
		if !cfg.Enabled {
			c.Next()
			return
		}

		// 检查是否在排除路径中
		for _, path := range cfg.ExcludePaths {
			if strings.HasPrefix(c.Request.URL.Path, path) {
				c.Next()
				return
			}
		}

		// 验证请求完整性
		if err := verifyRequestIntegrity(c, cfg); err != nil {
			logger.Error(c.Request.Context(), "SM3请求完整性校验失败",
				zap.Error(err),
				zap.String("method", c.Request.Method),
				zap.String("path", c.Request.URL.Path),
			)
			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
				"error": "请求数据完整性校验失败",
			})
			return
		}

		// 如果需要验证响应完整性，创建自定义响应写入器
		var writer gin.ResponseWriter
		var bodyBuffer *bytes.Buffer
		if cfg.VerifyResponses {
			bodyBuffer = &bytes.Buffer{}
			writer = &sm3ResponseWriter{
				ResponseWriter: c.Writer,
				body:          bodyBuffer,
			}
			c.Writer = writer
		}

		// 处理请求
		c.Next()

		// 添加响应完整性哈希头
		if cfg.VerifyResponses && bodyBuffer != nil && bodyBuffer.Len() > 0 {
			responseHash := calculateSM3Hash(bodyBuffer.Bytes(), cfg.HashFormat)
			c.Header("X-SM3-Hash", responseHash)
		}
	}
}

// verifyRequestIntegrity 验证请求完整性
func verifyRequestIntegrity(c *gin.Context, cfg SM3IntegrityConfig) error {
	// 对于GET、DELETE等方法，可能没有请求体，跳过验证
	if c.Request.Method == "GET" || c.Request.Method == "DELETE" || 
	   c.Request.Method == "HEAD" || c.Request.Method == "OPTIONS" {
		return nil
	}

	// 读取原始请求体
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return fmt.Errorf("读取请求体失败: %w", err)
	}

	// 恢复请求体
	c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
	c.Request.ContentLength = int64(len(body))

	// 如果请求体为空，跳过验证
	if len(body) == 0 {
		return nil
	}

	// 获取请求头中的哈希值
	providedHash := c.GetHeader(cfg.HashHeader)
	if providedHash == "" {
		return fmt.Errorf("缺少SM3哈希头: %s", cfg.HashHeader)
	}

	// 计算请求体的SM3哈希值
	calculatedHash := calculateSM3Hash(body, cfg.HashFormat)

	// 比较哈希值（忽略大小写，因为十六进制字符串可能有大小写差异）
	if strings.ToLower(providedHash) != strings.ToLower(calculatedHash) {
		return fmt.Errorf("SM3哈希值不匹配: 提供的=%s, 计算的=%s", providedHash, calculatedHash)
	}

	logger.Info(c.Request.Context(), "SM3请求完整性校验成功",
		zap.String("method", c.Request.Method),
		zap.String("path", c.Request.URL.Path),
	)

	return nil
}

// calculateSM3Hash 计算SM3哈希值
func calculateSM3Hash(data []byte, format string) string {
	switch format {
	case "base64":
		// 返回Base64编码的SM3哈希
		hashBytes := crypto.SM3Hash(data)
		return base64.StdEncoding.EncodeToString(hashBytes)
	case "hex":
		// 返回十六进制编码的SM3哈希（默认）
		return crypto.SM3HashHex(data)
	default:
		// 默认返回十六进制编码
		return crypto.SM3HashHex(data)
	}
}